
"""
世界观构建评分
"""
from models import Novel
from pydantic import Field
from core.utils import BaseAnalyzer, prompt_dir


class WorldviewRating(BaseAnalyzer):

    def __init__(self):
        super().__init__(prompt_file=prompt_dir / 'rating/r3. 世界观构建评分.md')

    def invoke(
            self,
            novel: Novel = Field(description="小说")
    ):
        all_worldviews = {
            chapter.chapter_title: chapter.worldview for chapter in novel.chapters
        }

        worldview_problems = dict()
        for chapter in novel.chapters:
            if not isinstance(chapter.problems, list):
                worldview_problems[chapter.chapter_title] = chapter.problems
                continue
            worldview_problems[chapter.chapter_title] = []
            for problem in chapter.problems:
                if not isinstance(problem, dict):
                    worldview_problems[chapter.chapter_title].append(problem)
                elif 'dimension' not in problem:
                    worldview_problems[chapter.chapter_title].append(problem)
                elif problem['dimension'] == '世界观构建':
                    worldview_problems[chapter.chapter_title].append(problem)

        prompt = self.template_prompt.format(
            inputs={
                "topn": novel.topn,
                "worldview": all_worldviews,
                "worldview_problems": worldview_problems,
            },
            remove_template_variables=True
        )

        novel.worldview_rating = self.analyze(prompt)
        return novel.worldview_rating
